# jupyter中使用paddlex需要设置matplotlib
import matplotlib

# 设置使用0号GPU卡（如无GPU，执行此代码后仍然会使用CPU训练模型）
import os

import paddlex as pdx
from paddlex.cls import transforms

matplotlib.use('Agg')
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

transforms_list = ['RandomHorizontalFlip', 'RandomVerticalFlip', 'Normalize', 'ResizeByShort', 'CenterCrop',
                   'RandomRotate', 'RandomDistort', 'RandomCrop']


# transforms_dict = {
#     'RandomHorizontalFlip': 'transforms.RandomHorizontalFlip',  # # 水平翻转
#     'RandomVerticalFlip': 'transforms.RandomVerticalFlip', # 垂直翻转
#     'Normalize': 'transforms.Normalize',               # 对图像进行标准化(归一化)
#     'ResizeByShort': 'transforms.ResizeByShort',       # resize
#     'CenterCrop': 'transforms.CenterCrop',             # 以图像中心点扩散裁剪长宽为`crop_size`的正方形
#     'RandomRotate': 'transforms.RandomRotate',         # 旋转
#     'RandomDistort': 'transforms.RandomDistort',        # 随机像素内容变换
#     'RandomCrop': 'transforms.RandomCrop'               # 随机剪裁
# }

def add_transforms(data: dict):
    train_transforms_list = [transforms.Normalize(), ]
    eval_transforms_list = [transforms.Normalize(), ]
    for i in data.keys():
        if i in transforms_list:
            transform = getattr(transforms, i)
            item = transform(**data[i])
            train_transforms_list.append(item)
    if 'CenterCrop' not in data.keys():
        transform = getattr(transforms, 'CenterCrop')
        item = transform()
        eval_transforms_list.append(item)
    train_transforms = transforms.Compose(train_transforms_list)
    eval_transforms = transforms.Compose(eval_transforms_list)
    return train_transforms, eval_transforms


def set_ImageNet_dataset(data_dir, transforms_obj, isshuffle=False):
    file_list = os.path.join(data_dir, 'train_list.txt')
    label_list = os.path.join(data_dir, 'labels.txt')
    dataset = pdx.datasets.ImageNet(
        data_dir=data_dir,
        file_list=file_list,
        label_list=label_list,
        transforms=transforms_obj,
        shuffle=isshuffle)
    # eval_dataset = pdx.datasets.ImageNet(
    #     data_dir=data_dir,
    #     file_list=file_list,
    #     label_list=label_list,
    #     transforms=eval_transforms)
    return dataset
